/*******************************************************************************
* This class uses libav libraries to handle video files.
* 
*
* To compile this code, you must add the libs -lavformat, -lavcodec, -lswscale
* to g++ compiler.
* 
*
* Licensed under GPL Version 3 license (/license.txt).
* Author: Eder A. Perez (eder.perez@ice.ufjf.br)
*******************************************************************************/

#include "cvvideo_vfile.h"


/**
  Object constructor
*/
CVVideoFile::CVVideoFile(CVpixel_fmt_t output_pxl_fmt): width(0), height(0), fps(0.f), bpp(0), bitrate(0),
                            errorcode(CV_SUCCESS), buffer(NULL), input_opened(false), fmt_ctx(NULL),
                            codec_ctx(NULL), frame(NULL), frameConv(NULL), pxlfmt(output_pxl_fmt)
{
  
  // Initialize libavcodec, and register all codecs and formats
  av_register_all();
  
  avformat_network_init();
  
}


/**
  Object destructor
*/
CVVideoFile::~CVVideoFile() {

  if (buffer) free(buffer);
  if (frameConv) av_free(frameConv);
  if (frame) av_free(frame);
  //if (codec_ctx) avcodec_close(codec_ctx);
  if (fmt_ctx) avformat_close_input(&fmt_ctx);
  
}


/**
  Open a video file. The video codec is auto-detected.
  If there's another video file opened it is closed first.

  @param filename Path to the video file

  @return Returns true if video file is successfully opened. Otherwise
  false is returned and errorcode is set.
*/
bool CVVideoFile::open(const char* filename) {
  AVCodec *codec;
  enum PixelFormat dstFmt;

  // Close the current input, if it exists, and open a new one
  if (input_opened) {
    this->~CVVideoFile();
    
    // Initiate attributes
    width = height = 0;
    fps = 0.f;
    bpp = 0;
    bitrate = 0;
    errorcode = CV_SUCCESS;
    buffer = NULL;
    input_opened = false;
    fmt_ctx = NULL;
    codec_ctx  = NULL;
    frame = frameConv = NULL; 
  }
  
  // Open video file (autodetect file format, buffer size and format parameters)
  if (avformat_open_input(&this->fmt_ctx, filename, NULL, NULL) != 0) {
    errorcode = CVERROR_INVALID_FILE;
    return false;
  }

  // Retrieve stream information
  if (avformat_find_stream_info(this->fmt_ctx, NULL) < 0) {
    errorcode = CVERROR_INVALID_FILE;
    avformat_close_input(&fmt_ctx);
    return false;
  }

  // Find the first video stream
  this->v_strm = -1;
  for (int i = 0; i < this->fmt_ctx->nb_streams; i++)
    if (this->fmt_ctx->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO) {
      this->v_strm = i;
      break;
    }

  if (this->v_strm == -1) {
    errorcode = CVERROR_NO_VIDEO_DATA;
    avformat_close_input(&fmt_ctx);
    return false;
  }

  // Get a pointer to the codec context for the video stream
  this->codec_ctx = this->fmt_ctx->streams[this->v_strm]->codec;

  // Find the decoder for the video stream
  codec = avcodec_find_decoder(this->codec_ctx->codec_id);
  if (codec == NULL) {
    errorcode = CVERROR_CODEC_NOT_FOUND;
    avformat_close_input(&fmt_ctx);
    this->codec_ctx = NULL;
    return false; 
  }

  // Open codec
  if (avcodec_open2(this->codec_ctx, codec, NULL) < 0) {
    errorcode = CVERROR_CODEC_NOT_OPENED;
    avformat_close_input(&fmt_ctx);
    this->codec_ctx = NULL;
    return false; 
  }

  // Allocate video frame
  this->frame = avcodec_alloc_frame();
  if (this->frame == NULL) {
    errorcode = CVERROR_MEMORY_ALLOCATION;
    avformat_close_input(&fmt_ctx);
    this->codec_ctx = NULL;
    return false;
  }

  this->frameConv = avcodec_alloc_frame();
  if (this->frameConv == NULL) {
    av_free(this->frame);
    this->frame = NULL;
    errorcode = CVERROR_MEMORY_ALLOCATION;
    avformat_close_input(&fmt_ctx);
    this->codec_ctx = NULL;
    return false;
  }

  // Determine required buffer size and allocate buffer
  dstFmt = get_libav_pxl_fmt();
  
  this->buffsize = avpicture_get_size(dstFmt, this->codec_ctx->width, this->codec_ctx->height);
  this->buffer = (CVubyte_t*) malloc(this->buffsize);
  
  // Assign appropriate parts of buffer to image planes in frameRGB
  avpicture_fill((AVPicture*)frameConv, buffer, dstFmt, codec_ctx->width, codec_ctx->height);

  // Set attributes
  this->width   = codec_ctx->width;
  this->height  = codec_ctx->height;
  this->bpp     = get_bpp();
  this->bitrate = fmt_ctx->bit_rate;
  this->fps = (float) 1.0f/(float)av_q2d(fmt_ctx->streams[v_strm]->time_base);
  
  // Set pixel format convertion context
  convert_context(this->width, this->height);
  
  input_opened = true;
  errorcode    = CV_SUCCESS;
  return true;
}



/** Gets current frame from video files. Frames from video files are in RGB with
* 24bpp.
* @params:
*   data: array of bytes.
*
* @return: returns true in success, false otherwise.
*       
**/
bool CVVideoFile::getFrame(CVubyte_t* data) {
  int got_picture;
  AVPacket packet;
  
  if (!data) return false;

  if (av_read_frame(this->fmt_ctx, &packet) >= 0) {

    // Is this a packet from the video stream?
    if (packet.stream_index == v_strm) {
      
      // Decode video frame
      avcodec_decode_video2(this->codec_ctx, this->frame, &got_picture, &packet);
      
      // Did we get a video frame?
      if (got_picture) {
        
	      // Pixel format convertion
	      sws_scale(this->conv_ctx, frame->data, frame->linesize, 0, this->codec_ctx->height, frameConv->data, frameConv->linesize);

			  // Copy the frame to data
        memcpy(data, frameConv->data[0], sizeof(CVubyte_t)*this->height*this->frameConv->linesize[0]);
			    
      }

      // Free the packet that was allocated by av_read_frame
      av_free_packet(&packet);
    }
    
    return true;
  }

  return false;
}



/**
  Resize the input frame grabbed.
  
  @param width New width
  @param height New height
  
  @return True if scale is possible, false otherwise.
*/
bool CVVideoFile::resize(int width, int height) {
  
  if(!this->input_opened) return false;
  
  free(this->buffer);
  
  this->buffsize = avpicture_get_size(PIX_FMT_RGB24, width, height);
  this->buffer = (CVubyte_t*) malloc(this->buffsize);
  
  // Assign appropriate parts of buffer to image planes
  av_free(frameConv);
  this->frameConv = avcodec_alloc_frame();
  avpicture_fill((AVPicture*)frameConv, buffer, PIX_FMT_RGB24, width, height);
  
  convert_context(width, height);
  
  this->width  = width;
  this->height = height;
  
  return true;
}



/**
  Creates a buffer of video frame size. It can be used as parameter for
  getFrame() method.
  
  @param buffer Points to created buffer
  
  @return The number of CVubyte_t allocated. Returns 0 if error.
*/
int CVVideoFile::createBuffer(CVubyte_t **buffer) {
  if (!input_opened) return 0;
  
  *buffer = (CVubyte_t*) calloc(this->buffsize, sizeof(CVubyte_t));
  
  if (!(*buffer)) {
    errorcode = CVERROR_MEMORY_ALLOCATION;
    return 0;
  }
  
  return this->buffsize;
}



/**
  Creates a CVImage object to hold frames grabbed.
  
  @param img Points to created buffer
  
  @return A new CVImage object. NULL if error.
*/
/*CVImage* CVVideoFile::createCVImage() {
  if(!this->input_opened) return NULL;
  
}*/



/**
  Set the values of SwsContext in order to convert/scale input frame.
  
  @param pxlfmt Converted pixel format
  @param width Converted width
  @param height Converted height
  
  @return True if context is created, false otherwise.
*/
bool CVVideoFile::convert_context(int width, int height) {
  enum PixelFormat dstFmt = get_libav_pxl_fmt();

  this->conv_ctx = sws_getCachedContext(this->conv_ctx, this->codec_ctx->width, this->codec_ctx->height,
                                        this->codec_ctx->pix_fmt, width, height, dstFmt, SWS_BICUBIC,
                                        NULL, NULL, NULL);


  if (!this->conv_ctx) {
    errorcode = CVERROR_CONVERTION;
    return false;
  }
  
  this->width = width;
  this->height = height;
  return true;
}



int CVVideoFile::get_bpp() {

  switch (this->pxlfmt) {
    case CV_PXL_FMT_GRAY8:
      return 8;
    
    case CV_PXL_FMT_YUV422:
      return 16;

    case CV_PXL_FMT_RGB24:
    default:
      return 24;
  }
 
}



enum PixelFormat CVVideoFile::get_libav_pxl_fmt() {

  switch (this->pxlfmt) {
    case CV_PXL_FMT_GRAY8:
      return PIX_FMT_GRAY8;
    
    case CV_PXL_FMT_YUV422:
      return PIX_FMT_YUV422P;

    case CV_PXL_FMT_RGB24:
    default:
      return PIX_FMT_RGB24;
  }
 
}



